from BaseNet import *
from cv2 import cvtColor as _cvtColor, resize as _resize
from cv2 import (
    rectangle as _rectangle,
    putText as _putText,
    getTextSize as _getTextSize,
)
from cv2 import COLOR_RGB2GRAY as _COLOR_RGB2GRAY, INTER_AREA as _INTER_AREA
from cv2 import FONT_HERSHEY_SIMPLEX as _FONT_HERSHEY_COMPLEX
from numpy import ndarray as _ndarray
from cv_helper import findBorderContours, show_plt, binarization, tosquare
from torch import no_grad as _no_grad, Tensor as _Tensor
from torchvision.transforms import ToTensor as _ToTensor

__all__ = [
    "compareMod",
    "to28",
    "distinguish",
    "distinguish_withMark",
]


def compareMod(mod: BaseNet) -> None:
    """对于每个输入打印出结果与正确答案进行对比"""
    with _no_grad():
        mod.eval()
        for i, l in mod.MNISTloader():
            # if UNLOCK_DEVICE:
            #     i = i.to(DEVICE)
            out: _Tensor = mod(i)
            show_plt(i[0][0], f"True answer is {l[0]}, return {out.argmax()}")


def to28(pic: _ndarray) -> _Tensor:
    """图片转为28x28大小张量"""
    img = tosquare(pic)
    img = _resize(img, (28, 28), interpolation=_INTER_AREA)
    img = binarization(img)
    return _ToTensor()(img).view(1, 1, 28, 28)


def distinguish(mod: BaseNet, pic: _ndarray):
    """借助opencv分区块识别"""
    ret: list[_Tensor] = []
    if len(pic.shape) > 2 and pic.shape[2] != 1:
        pic = _cvtColor(pic, _COLOR_RGB2GRAY)
    pic = binarization(pic, pic.mean() > 127)
    for border in findBorderContours(pic):
        imgTemp = pic[border.slice()]
        imgTemp = to28(imgTemp)
        tenTemp: _Tensor = mod(imgTemp)
        ret.append(tenTemp.argmax())
    return ret


def distinguish_withMark(mod: BaseNet, pic: _ndarray) -> _ndarray:
    """借助opencv分区块识别\n
    顺便把识别结果标记在图上"""
    img = pic.copy()
    if len(pic.shape) > 2 and pic.shape[2] != 1:
        pic = _cvtColor(pic, _COLOR_RGB2GRAY)
    pic = binarization(pic, pic.mean() > 127)
    for border in findBorderContours(pic):
        tenTemp: _Tensor = mod(to28(pic[border.slice()]))
        ans = int(tenTemp.argmax())
        print(f"The return is {ans}")
        _rectangle(img=img, pt1=border.begin, pt2=border.end, color=(0, 255, 0))
        hei, _ = _getTextSize(
            text=str(ans), fontFace=_FONT_HERSHEY_COMPLEX, fontScale=1, thickness=2
        )
        bottomLeft = (border.begin[0], border.begin[1] + hei[0])
        _putText(
            img=img,
            text=str(ans),
            org=bottomLeft,
            fontFace=_FONT_HERSHEY_COMPLEX,
            fontScale=1,
            color=(255, 0, 0),
            thickness=2,
        )
    return img
